package com.system.neural.trainer;

import com.aliasi.classify.Classification;
import com.aliasi.classify.Classified;
import com.aliasi.classify.DynamicLMClassifier;
import com.aliasi.lm.NGramProcessLM;
import com.system.neural.config.ClassModelConfig;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.FileOutputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import com.aliasi.util.Files;
import org.springframework.beans.factory.annotation.Autowired;

/**
 * ClassName: NGramClassierTrainer
 * FileName: NGramClassierTrainer.java
 * Description: NGram模型训练
 * History:
 * 版本号 			作者 			日期       				简介
 * 1.0				znlccy		    2021/10/4 12:17		    create
 */
@Slf4j
@Data
public class NGramClassierTrainer implements Serializable {

    /**
     * 声明序列化
     */
    private static final long serialVersionUID = 3405730611041981165L;

    /**
     * ngram特征取值，当前字符和前几个字符相关，以模型训练时候的参数设置为准，默认设置为3，太大容易出现拟合
     */
    private int ngramSize = 3;

    /**
     * 分类模型配置
     */
    @Autowired
    private ClassModelConfig classModelConfig;

    /**
     * 基于NGram提取特征之后的分类器
     */
    private DynamicLMClassifier<NGramProcessLM> classifier;

    /**
     * 训练模型
     * @throws Exception
     */
    public void trainModel() throws Exception {
        classifier = DynamicLMClassifier.createNGramProcess(this.getClassModelConfig().getCategories(), this.getNgramSize());
        for (int i = 0; i < this.getClassModelConfig().getCategories().length; i++) {
            File classDir = new File(this.getClassModelConfig().getTrainRootDir(), this.classModelConfig.getCategories()[i]);
            if (!classDir.isDirectory()) {
                String msg = "无法找到包含训练数据的路径 = " + classDir + "\n您是否未在训练root目录建立 " + this.getClassModelConfig().getCategories().length + "类别";
                log.error("错误信息:{}", msg);
                throw new IllegalArgumentException(msg);
            }

            String[] trainingFiles = classDir.list();
            for (int j = 0; j < trainingFiles.length; ++j) {
                File file = new File(classDir, trainingFiles[j]);
                String text = Files.readFromFile(file, "GBK");
                log.info("正在训练   -> " + this.getClassModelConfig().getCategories()[i] + "/" + trainingFiles[j]);
                Classification classification
                        = new Classification(this.getClassModelConfig().getCategories()[i]);
                Classified<CharSequence> classified
                        = new Classified<CharSequence>(text, classification);
                classifier.handle(classified);
            }
        }
        File modelFile = new File(this.getClassModelConfig().getTrainModelFile());
        System.out.println("所有类别均训练完成.开始保存模型到 " + modelFile.getAbsolutePath());
        if (modelFile.exists() == false) {
            log.info("指定模型文件不存在，将自动建立上层目录，并建立文件...");
            modelFile.getParentFile().mkdirs();
        } else {
            log.warn("指定模型文件已经存在，将覆盖...");
        }

        ObjectOutputStream os = new ObjectOutputStream(new FileOutputStream(
                modelFile));
        classifier.compileTo(os);
        os.close();

        log.info("模型保持成功！");
    }

    /*public static void main(String[] args) throws Exception {
        NGramClassierTrainer nGramClassierTrainer = new NGramClassierTrainer();
        nGramClassierTrainer.trainModel();
    }*/
}
